/*******************************************************************************
 * CopyRight (c) 2020 
 *
 * \file    mp4.c
 * \brief   Mp4 file format parser
 * \author  ZhangJing
 * \version 1.0.0
 * \date    2020-06-22
 *
 * -----------------------------------------------------------------------------
 * Modify history：
 * <Date>       | <Version> | <Author>  | <Discribe>
 * 2020-06-22   | v1.0.0    | ZhangJing | Create file
 * -----------------------------------------------------------------------------
 ******************************************************************************/

/*================================= include ==================================*/
#include <stdio.h>
#include <string.h>
#include <stdlib.h>

#include "publicdef.h"
#include "print_ctrl.h"
#include "format.h"
#include "mp4.h"

/*================================= struct ===================================*/
typedef struct Mp4BoxHeader {
    unsigned int type;
    long int len; /* total box length (excluding the length and type fields) */
} Mp4BoxHeader;

typedef struct Mp4ParserTable {
    unsigned int type;
    int (*parser)(void *, Mp4BoxHeader);
} Mp4ParserTable;

/*================================ variable ==================================*/
static int mp4_read_parentbox(void *format_context, Mp4BoxHeader atom);
static int mp4_read_trak(void *format_context, Mp4BoxHeader atom);
static int mp4_read_stco(void *format_context, Mp4BoxHeader atom);
static int mp4_read_stsz(void *format_context, Mp4BoxHeader atom);
static int mp4_read_co64(void *format_context, Mp4BoxHeader atom);
static int mp4_read_stsc(void *format_context, Mp4BoxHeader atom);
static int mp4_read_stsd(void *format_context, Mp4BoxHeader atom);

Mp4ParserTable mp4_parser_table[] = {
//{ MKTAG('f','t','y','p'), mp4_read_ftyp },
{ MKTAG('m','o','o','v'), mp4_read_parentbox },
//{ MKTAG('m','v','h','d'), mp4_read_mvhd },
{ MKTAG('t','r','a','k'), mp4_read_trak }, /**< relloc stream */
//{ MKTAG('t','k','h','d'), mp4_read_tkhd }, /* track header */
{ MKTAG('m','d','i','a'), mp4_read_parentbox },
//{ MKTAG('m','d','h','d'), mp4_read_mdhd },
//{ MKTAG('h','d','l','r'), mp4_read_hdlr },
{ MKTAG('m','i','n','f'), mp4_read_parentbox },
////{ MKTAG('v','m','h','d'), mp4_read_parentbox },
{ MKTAG('d','i','n','f'), mp4_read_parentbox },
//--{ MKTAG('d','r','e','f'), mp4_read_dref },
//{ MKTAG(0xa9,'u','r','l'), mp4_read_parentbox },
{ MKTAG('s','t','b','l'), mp4_read_parentbox },
{ MKTAG('s','t','s','d'), mp4_read_stsd }, /* sample description */
//--{ MKTAG('a','v','c','1'), mp4_read_avc1 },
//--{ MKTAG('h','v','c','1'), mp4_read_hvc1 },
//--{ MKTAG('h','v','c','C'), mp4_read_hvcC },
//{ MKTAG('a','v','c','C'), mp4_read_glbl },
//{ MKTAG('b','t','r','t'), mp4_read_parentbox },
//{ MKTAG('s','t','t','s'), mp4_read_stts },
//{ MKTAG('s','t','s','s'), mp4_read_stss }, /* sync sample */
{ MKTAG('s','t','s','c'), mp4_read_stsc },
{ MKTAG('s','t','s','z'), mp4_read_stsz }, /* sample size */
{ MKTAG('s','t','c','o'), mp4_read_stco }, /* sample offset */
{ MKTAG('c','o','6','4'), mp4_read_co64 }, /* sample offset */
{ MKTAG('u','d','t','a'), mp4_read_parentbox },
//--{ MKTAG('m','e','t','a'), mp4_read_parentbox },
//--{ MKTAG('i','l','s','t'), mp4_read_ilst },
////{ MKTAG('s','m','h','d'), mp4_read_parentbox },
//{ MKTAG('m','p','4','a'), mp4_read_parentbox },
//{ MKTAG('e','s','d','s'), mp4_read_esds },
{0, NULL}
};

/*================================ function ==================================*/
/*!
 * \brief   mp4 file format probe
 *
 * \param[in] format_context pointer to FormatContext
 * \return    0 if succeeded otherwise failed
 */
static int mp4_probe(void *format_context)
{
    FormatContext *fc = format_context;
    int ret = -EPROTONOSUPPORT;
    unsigned int tag;

    pr_entry(in);

    /** skip box length，get box type */
    memcpy(&tag, fc->probe_data + 4, sizeof(tag));

    //FIXME: checkout more tag
    switch (tag) {
    case MKTAG('m','o','o','v'): // metadata box
    case MKTAG('m','d','a','t'): // media data box
    case MKTAG('f','t','y','p'): // file type box
        ret = 0;
        break;
    default:
        break;
    }

    return ret;
}

/*!
 * \brief read parent box and child box with loop
 */
static int mp4_read_parentbox(void *format_context, Mp4BoxHeader atom)
{
    FormatContext *fc = format_context;
    IOContext *io = fc->io;
    int i;
    unsigned int tmp;
    long tmp64;
    long total_size = 0;
    Mp4BoxHeader a;

    while (total_size <= (atom.len - 8)) {
        int (*parser)(void *, Mp4BoxHeader) = NULL;
        if (atom.len >=8) {
            /** get child box */
            io->io_read(io, &tmp, 4);
            a.len = MKLE32(tmp);
            io->io_read(io, &tmp, 4);
            a.type = tmp;
            pr_detail("box size:%ld, name:%c%c%c%c", a.len, (char)a.type,
                (char)(a.type>>8), (char)(a.type>>16), (char)(a.type>>24));

            total_size += 8;
            if (a.len == 1 && total_size + 8 <= atom.len) {
                /* if 64 bit extended size */
                io->io_read(io, &tmp64, 8);
                a.len = MKLE64(tmp64) - 8;
                total_size += 8;
                pr_info("The full box len:%ld", a.len);
            }
            a.len -= 8;
            if (a.len < 0)
                break;
            a.len = MIN(a.len, atom.len - total_size);

            /* parser function */
            for (i = 0; mp4_parser_table[i].type; i++)
                if (mp4_parser_table[i].type == a.type) {
                    parser = mp4_parser_table[i].parser;
                    break;
                }

            /** parser a box */
            if (!parser) {
                /* if skip leaf atoms data */
                io->io_seek(io, a.len, SEEK_CUR);
            } else {
                parser(fc, a);
                //TODO: skip garbage at atom end
            }

            total_size += a.len;
        } else {
            pr_err("parent box length error:%d", atom.len);
            return -ESPIPE;
        }
    }
}

/*!
 * \brief read trak box and alloc stream struct
 */
static int mp4_read_trak(void *format_context, Mp4BoxHeader atom)
{
    FormatContext *fc = format_context;
    IOContext *io = fc->io;
    Stream *st;
    Stream **streams;
    Mp4Stream *ms;

    if ((fc->nb_streams <= 0) || !fc->streams) {
        streams = malloc(sizeof(*streams));
        if (!streams)
            return -ENOMEM;
    } else {
        streams = fc->streams;
        streams = realloc(streams, (fc->nb_streams + 1) * sizeof(*streams));
        if (!streams)
            return -ENOMEM;
    }
    fc->streams = streams;

    st = malloc(sizeof(Stream));
    if (!st)
        return -ENOMEM;
    st->index = fc->nb_streams;
    fc->streams[fc->nb_streams++] = st;

    ms = malloc(sizeof(Mp4Stream));
    if (!ms)
        return -ENOMEM;
    st->mp4st = ms;

    return mp4_read_parentbox(fc, atom);
}

/*!
 * \brief get all offset address of packets
 */
static int mp4_read_stco(void *format_context, Mp4BoxHeader atom)
{
    FormatContext *fc = format_context;
    IOContext *io = fc->io;
    Stream *st;
    Mp4Stream *ms;
    unsigned int i, tmp, entries;

    if (fc->nb_streams < 1)
        return -EMEDIUMTYPE; //FIXME: should skip the data len

    st = fc->streams[fc->nb_streams - 1];
    ms = (Mp4Stream *)st->mp4st;

    /* skip the 1byte version and 3bytes flags */
    io->io_seek(io, 4, SEEK_CUR);

    /* get packets count */
    io->io_read(io, &tmp, 4);
    entries = MKLE32(tmp);
    pr_debug("entries:%d", entries);

    if (!entries)
        return 0;

    ms->chunk_count = 0;
    ms->chunk_offsets = malloc(entries * sizeof(*ms->chunk_offsets));

    /* get packets offset */
    for (i = 0; i < entries; i++) {
        io->io_read(io, &tmp, 4);
        ms->chunk_offsets[i] = MKLE32(tmp);
        pr_detail("offsets[%d]:%d", i, ms->chunk_offsets[i]);
    }
    ms->chunk_count = i;

    return 0;
}

/*!
 * \brief get all offset address of packets
 */
static int mp4_read_co64(void *format_context, Mp4BoxHeader atom)
{
    FormatContext *fc = format_context;
    IOContext *io = fc->io;
    Stream *st;
    Mp4Stream *ms;
    unsigned int i, tmp, entries;
    long tmp64;

    if (fc->nb_streams < 1)
        return -EMEDIUMTYPE; //FIXME: should skip the data len

    st = fc->streams[fc->nb_streams - 1];
    ms = (Mp4Stream *)st->mp4st;

    /* skip the 1byte version and 3bytes flags */
    io->io_seek(io, 4, SEEK_CUR);

    /* get packets count */
    io->io_read(io, &tmp, 4);
    entries = MKLE32(tmp);
    pr_debug("entries:%d", entries);

    if (!entries)
        return 0;

    ms->chunk_count = 0;
    ms->chunk_offsets = malloc(entries * sizeof(*ms->chunk_offsets));

    /* get packets offset */
    for (i = 0; i < entries; i++) {
        io->io_read(io, &tmp64, 8);
        ms->chunk_offsets[i] = MKLE64(tmp64);
        pr_detail("offsets[%d]:%d", i, ms->chunk_offsets[i]);
    }
    ms->chunk_count = i;

    return 0;
}

/*!
 * \brief get all size of packets
 */
static int mp4_read_stsz(void *format_context, Mp4BoxHeader atom)
{
    FormatContext *fc = format_context;
    IOContext *io = fc->io;
    Stream *st;
    Mp4Stream *ms;
    unsigned int i, tmp, sample_size, sample_count;

    if (fc->nb_streams < 1)
        return -EMEDIUMTYPE; //FIXME: should skip the data len

    st = fc->streams[fc->nb_streams - 1];
    ms = (Mp4Stream *)st->mp4st;

    /* skip the 1byte version and 3bytes flags */
    io->io_seek(io, 4, SEEK_CUR);

    /* get packets count */
    io->io_read(io, &tmp, 4);
    sample_size = MKLE32(tmp);
    pr_detail("sample_size:%d", sample_size);
    io->io_read(io, &tmp, 4);
    sample_count = MKLE32(tmp);
    pr_detail("sample_count:%d", sample_count);

    if (!sample_count)
        return 0;

    ms->sample_count = 0;
    ms->sample_sizes = malloc(sample_count * sizeof(*ms->sample_sizes));

    for (i = 0; i < sample_count; i++) {
        io->io_read(io, &tmp, 4);
        ms->sample_sizes[i] = MKLE32(tmp);
        pr_detail("sample_sizes[%d]:%d", i, ms->sample_sizes[i]);
    }
    ms->sample_count = i;

    return 0;
}

/*!
 * \brief get sample describe with chunk
 */
static int mp4_read_stsc(void *format_context, Mp4BoxHeader atom)
{
    FormatContext *fc = format_context;
    IOContext *io = fc->io;
    Stream *st;
    Mp4Stream *ms;
    unsigned int i, tmp, stsc_count;
    BoxStsc *stsc_data;

    if (fc->nb_streams < 1)
        return -EMEDIUMTYPE; //FIXME: should skip the data len

    st = fc->streams[fc->nb_streams - 1];
    ms = (Mp4Stream *)st->mp4st;

    /* skip the 1byte version and 3bytes flags */
    io->io_seek(io, 4, SEEK_CUR);

    /* get sample-to-chunk count */
    io->io_read(io, &tmp, 4);
    stsc_count = MKLE32(tmp);
    pr_debug("stsc_count:%d", stsc_count);

    if (!stsc_count)
        return 0;

    ms->stsc_count = 0;
    ms->stsc_data = malloc(stsc_count * sizeof(*ms->stsc_data));

    for (i = 0; i < stsc_count; i++) {
        stsc_data = &ms->stsc_data[i];
        io->io_read(io, &tmp, 4);
        stsc_data->first_chunk = MKLE32(tmp);
        io->io_read(io, &tmp, 4);
        stsc_data->count_per_chunk = MKLE32(tmp);
        io->io_read(io, &tmp, 4);
        stsc_data->sample_id = MKLE32(tmp);
        pr_debug("stsc, firstChunk:%d samples:%d id:%d",
            ms->stsc_data[i].first_chunk, ms->stsc_data[i].count_per_chunk,
            ms->stsc_data[i].sample_id);
    }
    ms->stsc_count = i;

    return 0;
}

/*!
 * \brief get pps data
 */
static int mp4_read_stsd(void *format_context, Mp4BoxHeader atom)
{
    FormatContext *fc = format_context;
    IOContext *io = fc->io;
    Stream *st;
    Stream *st2;
    Mp4Stream *ms;

    unsigned int tag;
    unsigned char tag_size;
    char *tagstr = (char *)&tag;
    int total_len = atom.len;
    int cur_len = 0;
    unsigned int start_code = 0x01000000; /**< start code of NALU */

    if (fc->nb_streams < 1)
        goto exit;

    st = fc->streams[fc->nb_streams - 1];
    ms = st->mp4st;

    //FIXME: use better stsd parser
    /* skip 11 bytes of "stsd" */
    io->io_seek(io, 11, SEEK_CUR); cur_len += 11;

    io->io_read(io, &tag_size, 1); cur_len += 1;
    io->io_read(io, &tag, 4); cur_len += 4;
    pr_debug("TAG: %c%c%c%c, len: %d", tagstr[0], tagstr[1], tagstr[2], tagstr[3], tag_size);
    if (tag == MKTAG('a','v','c','1')) {
        st->streamname = "h264";
        st->id = ID_H264;
    } else if (tag == MKTAG('h','v','c','1')) {
        st->streamname = "hevc";
        st->id = ID_HEVC;
    } else if (tag == MKTAG('m','p','4','a')) {
        //BUG: should probe stream type
        if (fc->nb_streams > 1) {
            st2 = fc->streams[0];
            if (st2->id == ID_H264) {
                st->streamname = "mp3";
                st->id = ID_MP3;
            } else {
                st->streamname = "aac";
                st->id = ID_AAC;
            }
        }
    }

    /* malloc pps data */
    ms->pps_data = malloc(128);
    if (!ms->pps_data)
        goto exit;

    ms->pps_data_len = 0;
    if (tag == MKTAG('h','v','c','1')) {
        /* skip 81 bytes of "avc1" or "hvc1" */
        io->io_seek(io, 81, SEEK_CUR); cur_len += 81;

        io->io_read(io, &tag_size, 1); cur_len += 1;
        io->io_read(io, &tag, 4); cur_len += 4;
        pr_debug("TAG: %c%c%c%c, len: %d", tagstr[0], tagstr[1], tagstr[2], tagstr[3], tag_size);

        /* skip 23 bytes of "hvcC" */
        io->io_seek(io, 23, SEEK_CUR); cur_len += 23;

        /* skip 4 bytes */
        io->io_seek(io, 4, SEEK_CUR); cur_len += 4;
        io->io_read(io, &tag_size, 1); cur_len += 1;
        pr_debug("1 pps len: %d", tag_size);
        memcpy(ms->pps_data, &start_code, 4); ms->pps_data_len += 4;
        io->io_read(io, ms->pps_data + ms->pps_data_len, tag_size); cur_len += tag_size; ms->pps_data_len += tag_size;

        /* skip 4 bytes */
        io->io_seek(io, 4, SEEK_CUR); cur_len += 4;
        io->io_read(io, &tag_size, 1); cur_len += 1;
        pr_debug("2 pps len: %d", tag_size);
        memcpy(ms->pps_data + ms->pps_data_len, &start_code, 4); ms->pps_data_len += 4;
        io->io_read(io, ms->pps_data + ms->pps_data_len, tag_size); cur_len += tag_size; ms->pps_data_len += tag_size;

        /* skip 4 bytes */
        io->io_seek(io, 4, SEEK_CUR); cur_len += 4;
        io->io_read(io, &tag_size, 1); cur_len += 1;
        pr_debug("3 pps len: %d", tag_size);
        memcpy(ms->pps_data + ms->pps_data_len, &start_code, 4); ms->pps_data_len += 4;
        io->io_read(io, ms->pps_data + ms->pps_data_len, tag_size); cur_len += tag_size; ms->pps_data_len += tag_size;
    } else if (tag == MKTAG('a','v','c','1')) {
#define MKLE16(a) (((a)&0xFF00)>>8 | ((a)&0x00FF)<<8)
unsigned int tmp;
unsigned int sps_num, pps_num;
unsigned short tmp16, tag_len;
int i;
        /* skip 86 - 8 bytes of "avc1" or "hvc1" */
        io->io_seek(io, 78, SEEK_CUR); cur_len += 78;

        /* read size of the sub tag of hvc1 or avc1 */
        io->io_read(io, &tmp, 4); cur_len += 4;
        tmp = MKLE32(tmp);
        /* read the sub tag of hvc1 or avc1 */
        io->io_read(io, &tag, 4); cur_len += 4;
        pr_debug("TAG: %c%c%c%c, len: %d", tagstr[0], tagstr[1], tagstr[2], tagstr[3], tmp);

        /* get sps private data of avcC */
        if ((tag == MKTAG('a','v','c','C'))) {
            /* skip 5 bytes of version, ProfileIndi, compatibility, LevelIndi and NALU len */
            io->io_seek(io, 5, SEEK_CUR); cur_len += 5;

            /* get 1 byte of SPS num */
            io->io_read(io, &sps_num, 1); cur_len += 1;
            sps_num &= 0x1F;
            pr_debug("SPS num: %d", sps_num);
            for (i = 0; i < sps_num; i++) {
                /* read SPS len */
                io->io_read(io, &tmp16, 2); cur_len += 2;
                tag_len = MKLE16(tmp16);
                pr_debug("SPS len: %d", tag_len);
                /* read nal data */ //!!!: not to overflow the metadata buffer
                memcpy(ms->pps_data + ms->pps_data_len, &start_code, 4);
                ms->pps_data_len += 4;
                io->io_read(io, ms->pps_data + ms->pps_data_len, tag_len); cur_len += tag_len;
                ms->pps_data_len += tag_len;
            }

            /* get PPS */
            io->io_read(io, &pps_num, 1); cur_len += 1;
            pr_debug("PPS num: %d", pps_num);
            for (i = 0; i < pps_num; i++) {
                /* read PPS len */
                io->io_read(io, &tmp16, 2); cur_len += 2;
                tag_len = MKLE16(tmp16);
                pr_debug("PPS len: %d", tag_len);
                /* read nal data */ //!!!: not to overflow the metadata buffer
                memcpy(ms->pps_data + ms->pps_data_len, &start_code, 4);
                ms->pps_data_len += 4;
                io->io_read(io, ms->pps_data + ms->pps_data_len, tag_len); cur_len += tag_len;
                ms->pps_data_len += tag_len;
            }
            pr_debug("metadata len: %d", ms->pps_data_len);
        }
    }

exit:
    /* skip the bytes remaining */
    io->io_seek(io, total_len - cur_len, SEEK_CUR); cur_len += 4;

    return 0;
}

/*!
 * \brief     get stream information
 *
 * \param[in] format_context file data read
 * \return    0 if succeeded otherwise failed
 */
static int mp4_read_header(void *format_context)
{
    FormatContext *fc = format_context;
    IOContext *io = fc->io;

    Mp4BoxHeader atom = {MKTAG('r','o','o','t')};

    io->filesize = io->io_get_len(io);
    atom.len = io->filesize;
    pr_info("file len:%ld", atom.len);

    fc->nb_streams = 0;
    mp4_read_parentbox(fc, atom);

    return 0;
}

/*!
 * \brief     get stream data
 *
 * \param[in] format_context file data read
 * \return    0 if succeeded otherwise failed
 */
static int mp4_read_packet(void *format_context)
{
    //VOID
}

Format format_mp4 = {
    .formatname = "mp4",
    .id = ID_MP4,
    .read_probe = mp4_probe,
    .read_header = mp4_read_header,
    .read_packet = mp4_read_packet
};
